#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
企业知识检索智能体
Knowledge Retrieval Agent

基于向量检索的企业知识问答系统，支持：
- 公司政策查询（报销、请假、考勤等）
- 技术文档检索（API、配置、故障排查等）
- FAQ问答（常见问题解答）
- 产品文档查询（功能说明、使用教程等）
"""

import os
import logging
from typing import Dict, List, Any, Optional, Tuple
import dotenv

from langchain_community.embeddings import OpenAIEmbeddings
from pymilvus import MilvusClient

# 加载环境变量
dotenv.load_dotenv()

logger = logging.getLogger(__name__)

class KnowledgeAgent:
    """企业知识检索智能体"""
    
    def __init__(self):
        """初始化知识检索智能体"""
        try:
            # 初始化OpenAI配置
            self.openai_api_key = os.getenv("OPENAI_API_KEY")
            self.openai_base_url = os.getenv("OPENAI_BASE_URL")
            self.openai_model = os.getenv("OPENAI_MODEL", "4o-mini")
            self.openai_embedding_model = os.getenv("OPENAI_EMBEDDING_MODEL")
            
            # 初始化Milvus配置
            self.milvus_host = os.getenv("MILVUS_HOST")
            self.milvus_port = os.getenv("MILVUS_PORT")
            
            # 验证必要的环境变量
            if not all([self.openai_api_key, self.openai_base_url, self.openai_embedding_model, 
                       self.milvus_host, self.milvus_port]):
                raise ValueError("缺少必要的环境变量配置")
            
            # 初始化嵌入函数
            self.embedding_fn = OpenAIEmbeddings(
                api_key=self.openai_api_key,
                base_url=self.openai_base_url,
                model=self.openai_embedding_model
            )
            logger.info("知识检索智能体：嵌入函数初始化成功")
            
            # 初始化Milvus客户端
            self.milvus_client = MilvusClient(
                host=self.milvus_host,
                port=int(self.milvus_port)
            )
            logger.info("知识检索智能体：Milvus客户端初始化成功")
            
            # 暂时不使用LLM生成，直接基于检索结果回答
            logger.info("知识检索智能体：使用基于检索的回答模式")
            
            # 知识库配置
            self.collection_name = "knowledge_base"
            self.default_top_k = 3
            
            # 预加载集合
            self._preload_collection()
            
        except Exception as e:
            logger.error(f"知识检索智能体初始化失败: {e}")
            raise
    
    def _preload_collection(self):
        """预加载知识库集合"""
        try:
            self.milvus_client.load_collection(self.collection_name)
            logger.info(f"知识检索智能体：成功预加载集合 {self.collection_name}")
        except Exception as e:
            logger.warning(f"知识检索智能体：预加载集合失败 {e}")
    
    def _retrieve_knowledge(self, query: str, top_k: int = None) -> List[Dict[str, Any]]:
        """
        从知识库检索相关文档
        
        Args:
            query: 查询文本
            top_k: 返回结果数量
            
        Returns:
            检索结果列表
        """
        if top_k is None:
            top_k = self.default_top_k
            
        try:
            # 生成查询向量
            query_embedding = self.embedding_fn.embed_query(query.strip())
            
            # 执行向量检索
            results = self.milvus_client.search(
                collection_name=self.collection_name,
                data=[query_embedding],
                limit=top_k,
                output_fields=["category", "title", "source", "content"]
            )
            
            # 处理检索结果
            hits = results[0] if results else []
            knowledge_items = []
            
            for hit in hits:
                # 兼容不同版本的pymilvus返回格式
                fields = hit.get("entity", {}) if isinstance(hit, dict) else hit
                score = hit.get("distance") if isinstance(hit, dict) else getattr(hit, "distance", 0.0)
                
                # 提取字段信息
                category = fields.get("category") if isinstance(fields, dict) else getattr(fields, "category", "")
                title = fields.get("title") if isinstance(fields, dict) else getattr(fields, "title", "")
                source = fields.get("source") if isinstance(fields, dict) else getattr(fields, "source", "")
                content = fields.get("content") if isinstance(fields, dict) else getattr(fields, "content", "")
                
                knowledge_items.append({
                    "category": category,
                    "title": title,
                    "source": source,
                    "content": content,
                    "score": score
                })
            
            logger.info(f"知识检索成功：查询='{query}'，命中={len(knowledge_items)}条")
            return knowledge_items
            
        except Exception as e:
            logger.error(f"知识检索失败：{e}")
            return []
    
    def _generate_answer(self, query: str, knowledge_items: List[Dict[str, Any]]) -> str:
        """
        基于检索结果生成答案（简化版，直接基于检索内容）
        
        Args:
            query: 用户查询
            knowledge_items: 检索到的知识条目
            
        Returns:
            生成的答案
        """
        if not knowledge_items:
            return self._get_no_knowledge_response(query)
        
        try:
            # 获取最相关的文档
            best_match = knowledge_items[0]
            
            # 构建基于检索的回答
            answer_parts = [
                f"根据企业知识库，我找到了以下相关信息：\n",
                f"📄 **{best_match['title']}**",
                f"📂 类别：{best_match['category']}",
                f"\n📝 **详细内容：**",
                best_match['content']
            ]
            
            # 如果有多个相关结果，添加其他参考
            if len(knowledge_items) > 1:
                answer_parts.append(f"\n📚 **其他相关文档：**")
                for i, item in enumerate(knowledge_items[1:3], 2):  # 最多显示3个结果
                    answer_parts.append(f"• {item['title']} (相似度: {1-item['score']:.2f})")
            
            answer = "\n".join(answer_parts)
            
            logger.info(f"知识问答生成成功：查询='{query}'，命中={len(knowledge_items)}条")
            return answer
            
        except Exception as e:
            logger.error(f"答案生成失败：{e}")
            return f"抱歉，在处理知识库内容时遇到了问题：{str(e)[:100]}。请稍后重试。"
    
    def _get_no_knowledge_response(self, query: str) -> str:
        """当没有检索到相关知识时的回复"""
        return (
            f"抱歉，我没有在知识库中找到关于「{query}」的直接信息。\n\n"
            "您可以尝试：\n"
            "1. 使用更具体的关键词重新提问\n"
            "2. 描述更多背景信息或上下文\n"
            "3. 询问相关的政策、流程、技术文档或常见问题\n\n"
            "我可以帮您查询公司政策、技术文档、FAQ和产品使用指南等信息。"
        )
    
    def process_request(self, user_input: str, user_id: str, 
                       session_id: Optional[str] = None, 
                       context: Optional[dict] = None) -> dict:
        """
        处理知识检索请求
        
        Args:
            user_input: 用户输入
            user_id: 用户ID
            session_id: 会话ID（可选）
            context: 上下文信息（可选，包含查询重写信息）
            
        Returns:
            处理结果字典
        """
        try:
            logger.info(f"知识检索智能体处理请求：user_id={user_id}, query='{user_input}'")
            
            # 获取实际查询文本（优先使用重写后的查询）
            actual_query = user_input
            if context and context.get("rewritten_query") and context.get("rewrite_applied"):
                actual_query = context["rewritten_query"]
                logger.info(f"使用重写后查询：'{actual_query}'")
            
            # 执行知识检索
            knowledge_items = self._retrieve_knowledge(actual_query)
            
            # 生成答案
            answer = self._generate_answer(actual_query, knowledge_items)
            
            # 构建返回结果
            result = {
                "success": True,
                "output": answer,
                "error": None,
                "agent": "knowledge_agent",
                "meta": {
                    "original_query": user_input,
                    "actual_query": actual_query,
                    "knowledge_hits": len(knowledge_items),
                    "top_sources": [item["title"] for item in knowledge_items[:2]] if knowledge_items else []
                }
            }
            
            logger.info(f"知识检索智能体处理成功：命中={len(knowledge_items)}条知识")
            return result
            
        except Exception as e:
            error_msg = f"知识检索处理失败：{str(e)}"
            logger.error(error_msg)
            
            return {
                "success": False,
                "output": "抱歉，知识检索时出现了技术问题。请稍后重试，或者提供更多具体信息。",
                "error": error_msg,
                "agent": "knowledge_agent",
                "meta": {
                    "original_query": user_input,
                    "error_type": type(e).__name__
                }
            }


def create_knowledge_agent() -> KnowledgeAgent:
    """
    创建知识检索智能体实例
    
    Returns:
        KnowledgeAgent实例
    """
    try:
        return KnowledgeAgent()
    except Exception as e:
        logger.error(f"创建知识检索智能体失败：{e}")
        raise


def get_knowledge_agent() -> Optional[KnowledgeAgent]:
    """
    获取知识检索智能体实例（带错误处理）
    
    Returns:
        KnowledgeAgent实例或None
    """
    try:
        return create_knowledge_agent()
    except Exception as e:
        logger.error(f"获取知识检索智能体失败：{e}")
        return None


if __name__ == "__main__":
    # 测试代码
    agent = create_knowledge_agent()
    
    test_queries = [
        "报销多久到账",
        "API鉴权方式",
        "请假流程",
        "忘记打卡怎么办"
    ]
    
    for query in test_queries:
        print(f"\n测试查询: {query}")
        result = agent.process_request(query, "test_user")
        print(f"成功: {result['success']}")
        print(f"回答: {result['output'][:200]}...")
